//Examples/reed_sol_crs_01.cpp

#include <iostream>
#include <stdio.h>
#include <iostream>
#include <stdlib.h>
#include <string.h>
#include "jerasure.h"
#include "reed_sol.h"

#define talloc(type, num) (type *) malloc(sizeof(type)*(num))

void usage(char *s)
{
	fprintf(stderr, "usage: reed_sol_crs_01 k m w size - Does a simple Reed-Solomon coding example in GF(2^w).\n");
	fprintf(stderr, "       \n");
	fprintf(stderr, "       k+m must be <= 2^w.  w can be 8, 16 or 32.\n");
	fprintf(stderr, "       size must be a multiple of 8.\n");
	fprintf(stderr, "       It sets up a Cauchy distribution matrix and encodes\n");
	fprintf(stderr, "       k devices of 'size' bytes with it. Then it decodes.\n"); 
	fprintf(stderr, "       After that, it decodes device 0 by using JER_Slices::Make_Decoding_Matrix()\n");
	fprintf(stderr, "       and JER_Slices::Dotprod().\n");
	fprintf(stderr, "       \n");
	fprintf(stderr, "This demonstrates: CRS_Generator()\n");
	fprintf(stderr, "                   JER_Matrix::Print()\n");
	fprintf(stderr, "                   JER_Slices::Encode()\n");
	fprintf(stderr, "                   JER_Slices::Decode()\n");
	fprintf(stderr, "                   JER_Slices::Add_Drive_Failure()\n");
	fprintf(stderr, "                   JER_Slices::Make_Decoding_Matrix()\n");
	fprintf(stderr, "                   JER_Slices::Dotprod()\n");
	if (s != NULL) fprintf(stderr, "\n%s\n\n", s);
	exit(1);
}

void print_data_and_coding(JER_Slices *slices)
{
	int i, j, x;
	int n, sp;
	int k, m, w, size;
	uint64_t l;
	JER_Gen_T * g;
	vector < string > states;
	int max_data_state_len; //max strlen of a data drive state
	int max_coding_state_len;
	int state_len;

	states.resize(3);
	states[0] = "up";
	states[1] = "down";
	states[2] = "unusable";

	max_data_state_len = max_coding_state_len = 2;
	for(i=0;i<slices->N;i++){
		state_len = strlen(states[slices->States[i]].c_str());
		if(i<slices->K){
			//data
			if(state_len>max_data_state_len){
				max_data_state_len = state_len;
			}
		}else{
			//coding
			if(state_len>max_coding_state_len){
				max_coding_state_len = state_len;
			}
		}
	}

	g = slices->G;
	if(g == NULL){
		return;
	}

	if(g->M == NULL){
		return;
	}

	k = slices->K;
	m = slices->N - k;
	w = g->M->W;
	size = slices->PacketSize * slices->PacketsPerSlice;

	if (k > m) n = k;
	else n = m;
	sp = size * 2 + size/(w/8) + 8+max_data_state_len+3;

	printf("%-*sCoding\n", sp, "Data");
	for (i = 0; i < n; i++) {
		if (i < k) {
			printf("D%-2d (%-*s):", i, max_data_state_len,states[slices->States[i]].c_str());
			for (j = 0; j < size; j+=(w/8)) {
				printf(" ");
				for (x = 0; x < w/8; x++){
					printf("%02x", (unsigned char)slices->Ptrs[i][j+x]);
				}
			}
			printf("    ");
		}
		else printf("%*s", sp, "");
		if (i < m) {
			printf("C%-2d (%-*s):", i, max_coding_state_len,states[slices->States[i+k]].c_str());
			for (j = 0 ; j < size; j+=(w/8)) {
				printf(" ");
				for (x = 0; x < w/8; x++){
					printf("%02x", (unsigned char)slices->Ptrs[i+k][j+x]);
				}
			}
		}
		printf("\n");
	}
	printf("\n");
}

int main(int argc, char **argv)
{
	uint64_t l;
	int k, m, w, size;
	int i, j;
	int ps;
	JER_Matrix *jm;
	JER_Slices *slices;
	vector <int> erasures;
	vector <int> erased;
	JER_Matrix *dm;
	int *dm_ids;
	unsigned char *data;
	JER_Gen_T * g;
	vector <int> jm_super_row_ids;
	vector <int> dest_disk_ids;

	if (argc != 5) usage(NULL);
	if (sscanf(argv[1], "%d", &k) == 0 || k <= 0) usage((char *)"Bad k");
	if (sscanf(argv[2], "%d", &m) == 0 || m <= 0) usage((char *)"Bad m");
	if (sscanf(argv[3], "%d", &w) == 0 || (w != 8 && w != 16 && w != 32))
		usage((char *)"Bad w");
	if (w < 32 && k + m > (1 << w)) usage((char *)"k + m must be <= 2 ^ w");
	if (sscanf(argv[4], "%d", &size) == 0 || size%8 != 0) 
		usage((char *)"size must be multiple of 8");

	ps = sizeof(uint64_t);
	g = CRS_Generator(k,m,w);
	if(g == NULL){
		printf("Error creating generator\n");
		return 0;	
	}

	jm = g->M;
	printf("The Coding Matrix (the last m rows of the Distribution Matrix):\n\n");
	jm->Print();
	printf("\n");

	srand48(0);

	//allocate and create empty slices
	slices = new JER_Slices(k+m, k,ps, size/ps,g);
	
	data = new unsigned char [size * (k + m)];

	for (i = 0; i < k; i++) {
		for(j = 0; j < size; j+=ps) {
			l = lrand48();
			l <<= 8*4;
			l += lrand48();

			memcpy(data + i*size + j, &l, ps);
		}
		slices->Ptrs[i] = data + i*size;
	}

	for (i = 0; i < m; i++) {
		slices->Ptrs[k + i] = data + k*size + i*size;
	}

	printf("Original data:\n\n");
	print_data_and_coding(slices);

	slices->Encode();

	printf("Encoding complete:\n\n");
	print_data_and_coding(slices);

	//erase m drives
	erasures.resize(m,0);
	for (i = 0; i < m; ) {
		erasures[i] = lrand48()%(k+m);
		if (slices->States[erasures[i]] == 0) {
			slices->States[erasures[i]] = 1;
			bzero(slices->Ptrs[erasures[i]], size);
			i++;
		}
	}

	printf("Erased %d random devices:\n\n", m);
	print_data_and_coding(slices);

	if(slices->Decode() == -1){
		printf("Error decoding\n");
		return 0;
	}

	printf("State of the system after decoding:\n\n");
	print_data_and_coding(slices);

	dm_ids = new int [k];
	dm = new JER_Matrix(k, k, w);

	printf("Erased the first %d devices\n\n",m);
	for (i = 0; i < m; i++) {
		slices->Add_Drive_Failure(i);
		bzero(slices->Ptrs[i], size);
	}
	print_data_and_coding(slices);

	dm = slices->Make_Decoding_Matrix(dm_ids);
	printf("Here is the decoding matrix:\n\n", m);
	dm->Print();

	printf("\n");
	printf("And dm_ids:\n\n");
	for (i = 0; i < k; i++) {
		printf(" ");
		if (i % w == 0 && i > 0) printf(" ");
		printf("%d", dm_ids[i]);
	}
	printf("\n");

	bzero((char *)data, size);
	jm_super_row_ids.push_back(0);
	dest_disk_ids.push_back(0);
	slices->Dotprod(dm, jm_super_row_ids, dest_disk_ids,dm_ids);

	printf("\nAfter calling JER_Slices::Dotprod(), we calculate the value of device #0 to be:\n\n");
	printf("D0 :");
	for (i = 0; i < size; i+=(w/8)) {
		printf(" ");
		for (j = 0; j < w/8; j++){
			printf("%02x", (unsigned char)slices->Ptrs[0][i+j]);
		}
	}
	printf("\n\n");

	delete g;
	delete slices;
	delete dm;

	return 0;
}
